import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Union
import logging
from datetime import datetime, timedelta
import json
from dataclasses import dataclass
from abc import ABC, abstractmethod

@dataclass
class LLMResponse:
    """LLM response data structure"""
    content: str
    confidence: float
    metadata: Dict[str, Any]
    timestamp: datetime
    model_used: str
    tokens_used: int
    cost: float = 0.0

@dataclass
class TradingInsight:
    """Trading insight generated by LLM"""
    insight_type: str  # 'signal', 'analysis', 'recommendation', 'risk_warning'
    content: str
    confidence: float
    symbols: List[str]
    timeframe: str
    reasoning: str
    timestamp: datetime

class LLMProvider(ABC):
    """Abstract base class for LLM providers"""
    
    @abstractmethod
    def generate_response(self, prompt: str, **kwargs) -> LLMResponse:
        """Generate response from LLM"""
        pass
    
    @abstractmethod
    def get_model_info(self) -> Dict[str, Any]:
        """Get model information"""
        pass

class OpenAIProvider(LLMProvider):
    """OpenAI LLM provider"""
    
    def __init__(self, api_key: str, model: str = "gpt-3.5-turbo"):
        self.api_key = api_key
        self.model = model
        self.logger = logging.getLogger(__name__)
        
        try:
            import openai
            openai.api_key = api_key
            self.client = openai
        except ImportError:
            raise ImportError("OpenAI package not installed. Install with: pip install openai")
    
    def generate_response(self, prompt: str, **kwargs) -> LLMResponse:
        """Generate response using OpenAI"""
        try:
            response = self.client.ChatCompletion.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "You are a financial analyst and trading expert."},
                    {"role": "user", "content": prompt}
                ],
                temperature=kwargs.get('temperature', 0.3),
                max_tokens=kwargs.get('max_tokens', 1000)
            )
            
            return LLMResponse(
                content=response.choices[0].message.content,
                confidence=0.8,  # OpenAI doesn't provide confidence scores
                metadata={'finish_reason': response.choices[0].finish_reason},
                timestamp=datetime.now(),
                model_used=self.model,
                tokens_used=response.usage.total_tokens,
                cost=self._calculate_cost(response.usage)
            )
            
        except Exception as e:
            self.logger.error(f"OpenAI API error: {e}")
            raise
    
    def get_model_info(self) -> Dict[str, Any]:
        """Get OpenAI model information"""
        return {
            'provider': 'OpenAI',
            'model': self.model,
            'max_tokens': 4096 if 'gpt-3.5' in self.model else 8192,
            'supports_functions': True
        }
    
    def _calculate_cost(self, usage) -> float:
        """Calculate cost based on token usage"""
        # Approximate costs (may vary)
        costs = {
            'gpt-3.5-turbo': {'input': 0.0015, 'output': 0.002},
            'gpt-4': {'input': 0.03, 'output': 0.06}
        }
        
        model_costs = costs.get(self.model, costs['gpt-3.5-turbo'])
        return (usage.prompt_tokens * model_costs['input'] + 
                usage.completion_tokens * model_costs['output']) / 1000

class LocalLLMProvider(LLMProvider):
    """Local LLM provider using transformers"""
    
    def __init__(self, model_name: str = "microsoft/DialoGPT-medium"):
        self.model_name = model_name
        self.logger = logging.getLogger(__name__)
        
        try:
            from transformers import AutoTokenizer, AutoModelForCausalLM
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            self.model = AutoModelForCausalLM.from_pretrained(model_name)
            self.tokenizer.pad_token = self.tokenizer.eos_token
        except ImportError:
            raise ImportError("Transformers package not installed. Install with: pip install transformers torch")
    
    def generate_response(self, prompt: str, **kwargs) -> LLMResponse:
        """Generate response using local model"""
        try:
            inputs = self.tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=512)
            
            with torch.no_grad():
                outputs = self.model.generate(
                    inputs,
                    max_length=kwargs.get('max_length', 200),
                    temperature=kwargs.get('temperature', 0.7),
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id
                )
            
            response_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            response_text = response_text[len(prompt):].strip()
            
            return LLMResponse(
                content=response_text,
                confidence=0.6,  # Local models typically have lower confidence
                metadata={'model_name': self.model_name},
                timestamp=datetime.now(),
                model_used=self.model_name,
                tokens_used=len(outputs[0])
            )
            
        except Exception as e:
            self.logger.error(f"Local LLM error: {e}")
            raise
    
    def get_model_info(self) -> Dict[str, Any]:
        """Get local model information"""
        return {
            'provider': 'Local',
            'model': self.model_name,
            'max_tokens': 512,
            'supports_functions': False
        }

class LLMIntegration:
    """LangChain + LLM integration for trading system"""
    
    def __init__(self, provider: str = "openai", api_key: str = None, 
                 model: str = "gpt-3.5-turbo"):
        self.logger = logging.getLogger(__name__)
        self.provider = self._initialize_provider(provider, api_key, model)
        
        # Initialize LangChain components
        self._init_langchain()
        
        # Trading-specific prompts
        self.trading_prompts = {
            'market_analysis': self._get_market_analysis_prompt(),
            'signal_generation': self._get_signal_generation_prompt(),
            'risk_assessment': self._get_risk_assessment_prompt(),
            'portfolio_optimization': self._get_portfolio_optimization_prompt()
        }
    
    def _initialize_provider(self, provider: str, api_key: str, model: str) -> LLMProvider:
        """Initialize LLM provider"""
        if provider.lower() == "openai":
            if not api_key:
                raise ValueError("OpenAI API key required")
            return OpenAIProvider(api_key, model)
        elif provider.lower() == "local":
            return LocalLLMProvider(model)
        else:
            raise ValueError(f"Unsupported provider: {provider}")
    
    def _init_langchain(self):
        """Initialize LangChain components"""
        try:
            from langchain.llms import OpenAI
            from langchain.chains import LLMChain
            from langchain.prompts import PromptTemplate
            from langchain.memory import ConversationBufferMemory
            
            self.langchain_available = True
            self.memory = ConversationBufferMemory()
            
        except ImportError:
            self.logger.warning("LangChain not available. Install with: pip install langchain")
            self.langchain_available = False
    
    def analyze_market_data(self, market_data: pd.DataFrame, 
                           symbols: List[str]) -> TradingInsight:
        """Analyze market data and generate insights"""
        prompt = self._create_market_analysis_prompt(market_data, symbols)
        
        try:
            response = self.provider.generate_response(prompt)
            
            # Parse response to extract insights
            insight = self._parse_trading_insight(response.content, 'analysis', symbols)
            return insight
            
        except Exception as e:
            self.logger.error(f"Market analysis failed: {e}")
            return self._create_default_insight('analysis', symbols)
    
    def generate_trading_signals(self, factor_data: pd.DataFrame,
                                price_data: pd.DataFrame,
                                strategy_context: str = "") -> TradingInsight:
        """Generate trading signals based on factor and price data"""
        prompt = self._create_signal_generation_prompt(factor_data, price_data, strategy_context)
        
        try:
            response = self.provider.generate_response(prompt)
            insight = self._parse_trading_insight(response.content, 'signal', 
                                                list(factor_data.columns))
            return insight
            
        except Exception as e:
            self.logger.error(f"Signal generation failed: {e}")
            return self._create_default_insight('signal', list(factor_data.columns))
    
    def assess_risk(self, portfolio_data: Dict[str, Any],
                   market_conditions: Dict[str, Any]) -> TradingInsight:
        """Assess portfolio risk using LLM"""
        prompt = self._create_risk_assessment_prompt(portfolio_data, market_conditions)
        
        try:
            response = self.provider.generate_response(prompt)
            symbols = list(portfolio_data.get('positions', {}).keys())
            insight = self._parse_trading_insight(response.content, 'risk_warning', symbols)
            return insight
            
        except Exception as e:
            self.logger.error(f"Risk assessment failed: {e}")
            return self._create_default_insight('risk_warning', [])
    
    def optimize_portfolio(self, current_weights: Dict[str, float],
                          factor_scores: Dict[str, Dict[str, float]],
                          constraints: Dict[str, Any]) -> TradingInsight:
        """Generate portfolio optimization recommendations"""
        prompt = self._create_portfolio_optimization_prompt(current_weights, factor_scores, constraints)
        
        try:
            response = self.provider.generate_response(prompt)
            symbols = list(current_weights.keys())
            insight = self._parse_trading_insight(response.content, 'recommendation', symbols)
            return insight
            
        except Exception as e:
            self.logger.error(f"Portfolio optimization failed: {e}")
            return self._create_default_insight('recommendation', list(current_weights.keys()))
    
    def answer_trading_question(self, question: str, 
                               context_data: Dict[str, Any] = None) -> LLMResponse:
        """Answer trading-related questions"""
        prompt = self._create_question_prompt(question, context_data)
        
        try:
            response = self.provider.generate_response(prompt)
            return response
            
        except Exception as e:
            self.logger.error(f"Question answering failed: {e}")
            return LLMResponse(
                content="I'm unable to answer that question at the moment.",
                confidence=0.0,
                metadata={'error': str(e)},
                timestamp=datetime.now(),
                model_used=self.provider.get_model_info()['model'],
                tokens_used=0
            )
    
    def _create_market_analysis_prompt(self, market_data: pd.DataFrame, 
                                     symbols: List[str]) -> str:
        """Create market analysis prompt"""
        data_summary = market_data.describe().to_string()
        
        return f"""
        Analyze the following market data for symbols {symbols}:
        
        Market Data Summary:
        {data_summary}
        
        Please provide:
        1. Market trend analysis
        2. Key technical levels
        3. Potential catalysts
        4. Risk factors
        5. Trading opportunities
        
        Format your response as JSON with keys: trend, levels, catalysts, risks, opportunities
        """
    
    def _create_signal_generation_prompt(self, factor_data: pd.DataFrame,
                                       price_data: pd.DataFrame,
                                       strategy_context: str) -> str:
        """Create signal generation prompt"""
        factor_summary = factor_data.tail(5).to_string()
        price_summary = price_data.tail(5).to_string()
        
        return f"""
        Generate trading signals based on the following data:
        
        Factor Data (last 5 periods):
        {factor_summary}
        
        Price Data (last 5 periods):
        {price_summary}
        
        Strategy Context: {strategy_context}
        
        Provide trading signals in JSON format with:
        - signal_type: 'buy', 'sell', or 'hold'
        - confidence: 0-1
        - reasoning: explanation
        - target_price: if applicable
        - stop_loss: if applicable
        """
    
    def _create_risk_assessment_prompt(self, portfolio_data: Dict[str, Any],
                                     market_conditions: Dict[str, Any]) -> str:
        """Create risk assessment prompt"""
        return f"""
        Assess portfolio risk based on:
        
        Portfolio Data:
        {json.dumps(portfolio_data, indent=2)}
        
        Market Conditions:
        {json.dumps(market_conditions, indent=2)}
        
        Provide risk assessment in JSON format with:
        - overall_risk: 'low', 'medium', 'high'
        - risk_factors: list of specific risks
        - recommendations: risk mitigation suggestions
        - position_adjustments: suggested changes
        """
    
    def _create_portfolio_optimization_prompt(self, current_weights: Dict[str, float],
                                            factor_scores: Dict[str, Dict[str, float]],
                                            constraints: Dict[str, Any]) -> str:
        """Create portfolio optimization prompt"""
        return f"""
        Optimize portfolio weights based on:
        
        Current Weights:
        {json.dumps(current_weights, indent=2)}
        
        Factor Scores:
        {json.dumps(factor_scores, indent=2)}
        
        Constraints:
        {json.dumps(constraints, indent=2)}
        
        Provide optimization recommendations in JSON format with:
        - suggested_weights: new weight allocation
        - expected_return: estimated return
        - risk_level: portfolio risk assessment
        - rebalancing_actions: specific actions to take
        """
    
    def _create_question_prompt(self, question: str, context_data: Dict[str, Any] = None) -> str:
        """Create question answering prompt"""
        context = ""
        if context_data:
            context = f"\nContext: {json.dumps(context_data, indent=2)}"
        
        return f"""
        Answer the following trading question:
        
        Question: {question}
        {context}
        
        Provide a comprehensive, accurate answer based on financial knowledge and best practices.
        """
    
    def _parse_trading_insight(self, content: str, insight_type: str, 
                              symbols: List[str]) -> TradingInsight:
        """Parse LLM response into TradingInsight"""
        try:
            # Try to parse JSON response
            if content.strip().startswith('{'):
                data = json.loads(content)
                reasoning = data.get('reasoning', content)
            else:
                reasoning = content
            
            return TradingInsight(
                insight_type=insight_type,
                content=content,
                confidence=0.7,  # Default confidence
                symbols=symbols,
                timeframe='current',
                reasoning=reasoning,
                timestamp=datetime.now()
            )
            
        except json.JSONDecodeError:
            # If not JSON, use content as is
            return TradingInsight(
                insight_type=insight_type,
                content=content,
                confidence=0.5,
                symbols=symbols,
                timeframe='current',
                reasoning=content,
                timestamp=datetime.now()
            )
    
    def _create_default_insight(self, insight_type: str, symbols: List[str]) -> TradingInsight:
        """Create default insight when LLM fails"""
        return TradingInsight(
            insight_type=insight_type,
            content="Unable to generate insight due to technical issues.",
            confidence=0.0,
            symbols=symbols,
            timeframe='current',
            reasoning="LLM service unavailable",
            timestamp=datetime.now()
        )
    
    def _get_market_analysis_prompt(self) -> str:
        """Get market analysis system prompt"""
        return """You are an expert financial analyst specializing in market analysis. 
        Provide accurate, data-driven insights for trading decisions."""
    
    def _get_signal_generation_prompt(self) -> str:
        """Get signal generation system prompt"""
        return """You are a quantitative trading strategist. 
        Generate trading signals based on factor analysis and market data."""
    
    def _get_risk_assessment_prompt(self) -> str:
        """Get risk assessment system prompt"""
        return """You are a risk management expert. 
        Assess portfolio risks and provide mitigation strategies."""
    
    def _get_portfolio_optimization_prompt(self) -> str:
        """Get portfolio optimization system prompt"""
        return """You are a portfolio optimization specialist. 
        Provide optimal weight allocations based on factor analysis."""
    
    def get_provider_info(self) -> Dict[str, Any]:
        """Get information about the current LLM provider"""
        return self.provider.get_model_info()
    
    def get_usage_stats(self) -> Dict[str, Any]:
        """Get usage statistics"""
        return {
            'provider': self.provider.get_model_info()['provider'],
            'model': self.provider.get_model_info()['model'],
            'timestamp': datetime.now()
        } 